﻿using System;
using System.IO;
using System.Collections.Generic;
using System.Threading; 
using ASquared;
using ASquared.SymbolicMath;
using ASquared.ModelStatistics;
using Microsoft.SolverFoundation.Services;
using Microsoft.SolverFoundation.Common;
using Microsoft.SolverFoundation.Solvers;

namespace ASquared.ModelOptimization
{
    /// <summary>A conditional gradient solver.</summary>
    public class ConditionalGradientSolver : ILineSearchSolvable
    {
        #region Instance variables

        private IConditionalGradientSolvableModel model; // this serves as the linear optimization solver if provided
        private IConditionalGradientSolvableFunction objective;
        private Matrix xcurr = null, xprev0 = null, xprev1 = null, xprev2 = null;
        private Matrix xbar = null, xbarprev = null, xbarprev1 = null, xbarprev2 = null;
        private Matrix _gradient;
        private Symbol _alpha = 1;
        private Symbol _lsFxn, _currFxn;
        private Double _defLStolerance = 0.0001, _lstolerance = 0.0001;
        private Double _origtolerance = 0.00001, _tolerance = 0.00001;
        private Int32 _iter, _maxIters = 1000;
        private Int32 _iterToRunUnconstrained = -1; 
        private Double _totalOptimum = 0, 
            _optimum = 0,
            _prevOptimum = Double.NaN, 
            _prevOptimum1 = Double.NaN, 
            _prevOptimum2 = Double.NaN;
        private Double _lstoleranceUpdateFactor = 1.0;
        private string _logFile = "";
        private Boolean _triedBoundarySolution = false;
        private StreamWriter _sw = null;
        private String _ConvergedComponents;
        private EventHandler _initHandler, _iterBottomHandler, _convergedHandler;
        private Boolean _handleEvents;
        private Thread _thread = null;
        private Boolean _initialized; 

        public event EventHandler Init;
        public event EventHandler IterBottom;
        public event EventHandler Converged;

        #endregion

        #region Properties

        /// <summary>Gets and sets the current tolerance for the line search routine.</summary>
        /// <remarks>The line search tolerance is reset at the beginning of the solve method because the tolerance is updated during iterations of the conditional gradient method.</remarks>
        public Double LineSearchTolerance { get { return _lstolerance; } set { _defLStolerance = _lstolerance = value; } }
        /// <summary>Gets and sets the factor used to update the line search tolerance iteration.</summary>
        public Double LineSearchToleranceUpdateFactor { get { return _lstoleranceUpdateFactor; } set { _lstoleranceUpdateFactor = value; } }
        /// <summary>Iteration on which to run an unconstrained optimization problem in an attempt to speed it up.</summary>
        public Int32 IterationToRunUnconstrained { get { return _iterToRunUnconstrained; } set { _iterToRunUnconstrained = value; } }
        /// <summary>Gets and sets the tolerance of the function</summary>
        public Double Tolerance { get { return _tolerance; } set { _origtolerance = _tolerance = value; } }
        /// <summary>Gets and sets the maximum number of iterations</summary>
        public Int32 MaxIterations { get { return _maxIters; } set { _maxIters = value; } }
        /// <summary>Gets the iteration of the conditional gradient method</summary>
        public Int32 Iteration { get { return _iter; } }
        /// <summary>Gets the decisions made after solution was reached.</summary>
        public Matrix Decisions { get { return xbar; } }
        /// <summary>Gets the optimum value of the objective function at the current decisions.</summary>
        public Double Optimum { get { return _optimum; } }
        /// <summary>Gets the total optimum value of the objective function after all convergences</summary>
        public Double TotalOptimum { get { return _totalOptimum; } }
        /// <summary>Gets and sets whether or not to write the log file.</summary>
        public string LogFile { get { return _logFile; } set { _logFile = value; } }
        /// <summary>Gets and sets whether this solver has converged or not.</summary>
        public bool IsConverged { get { return this.objective.IsConverged; } set { this.objective.IsConverged = value; } }
        /// <summary>Specifies whether to try a boundary solution.</summary>
        public bool TrySpeedingIterationsWithBoundarySolution { get; set; } 

        #endregion

        #region Constructors

        public ConditionalGradientSolver(IConditionalGradientSolvableModel model)
        {
            this.model = model;
            this.objective = model.objective;
            _initHandler = new EventHandler(model_Init);
            _iterBottomHandler = new EventHandler(model_IterBottom);
            _convergedHandler = new EventHandler(model_Converged);
            this.model.Init += _initHandler;
            this.model.IterBottom += _iterBottomHandler;
            this.model.Converged += _convergedHandler;
            _handleEvents = true;
            this.TrySpeedingIterationsWithBoundarySolution = false;
        }

        public ConditionalGradientSolver(IConditionalGradientSolvableFunction fxn)
        {
            this.objective = fxn;
            this.model = null;
            this.TrySpeedingIterationsWithBoundarySolution = false;
        }

        #endregion

        #region Methods

        public void Solve()
        {
            Solve(true);
        }

        public void Solve(Boolean initialize)
        {
            if (this.model != null)
            {
                this.model.Run();
            }
            else
            {
                _triedBoundarySolution = false;
                if (initialize) 
                    this.objective.Initialize();
                for (_iter = 0; !this.objective.IsConverged && _iter < _maxIters; _iter++)
                    this.SolveIter(_iter);
                //if (_iter == _maxIters)
                //    throw new Exception("Maximum iterations were reached when attempting to solve. Solution not found.");
            }
            this.CloseLog();
        }
        public void Reset()
        {
            _iter = -1;
            _prevOptimum2 = Double.NaN;
            _prevOptimum1 = Double.NaN;
            _prevOptimum = Double.NaN;
            _triedBoundarySolution = false;
            this.xprev2 = null;
            this.xprev1 = null;
            this.xprev0 = null;
            this.xbarprev2 = null;
            this.xbarprev1 = null;
            this.xbarprev = null;
            this.model.Reset();
        }
        private void SavePrevious()
        {
            this.xprev2 = this.xprev1;
            this.xprev1 = this.xprev0;
            this.xprev0 = this.xcurr;
            this.xbarprev2 = this.xbarprev1;
            this.xbarprev1 = this.xbarprev;
            this.xbarprev = this.xbar;
        }
        private void SolveIter(Int32 iter)
        {
            if (this.model == null)
                this.SolveLinearOptimization();
            this.xcurr = this.objective.decisions;
            if (this.xbar != null)
                this.objective.decisions = this.xbar; // make sure that the objective function is up-to-date 
            _currFxn = this.objective.fxn;
            if (this.IterationToRunUnconstrained == _iter && false) // not running this because it throws off the solution
                this.RunUnconstrainedProblem();
            if (iter > 0)
            {
                this.LineSearch();
                if (this.model != null) this.objective.decisions = this.xbar;
            }
            else
            {
                this.xbar = this.xcurr;
                _optimum = (this.objective.costs.Mult(this.xbar, TransposeOption.LeftAsNeeded))[0];
            }
            this.DifferentiateFunction();
            this.UpdateConvergence();
            this.UpdateCosts();
            this.WriteToLog();
        }
        private void SolveLinearOptimization()
        {
            LinearSolver ls = new LinearSolver(this.objective.A, this.objective.b, this.objective.costs);
            ls.ConstraintTypes = this.objective.ConstraintTypes;
            ls.VariableTypes = this.objective.VariableTypes;
            ls.Solve();
            this.objective.decisions = ls.x;
        }
        private object[] LineSearchSubstitution()
        {
            object[] objs = new object[this.objective.variables.Length];
            Symbol a = 'a';
            if (this.xprev2 != null && !_triedBoundarySolution && this.xprev2 == this.xprev0 && this.xprev1 == this.xcurr && this.TrySpeedingIterationsWithBoundarySolution)
            {
                _triedBoundarySolution = true;
                this.xbar = this.xprev0;
            }
            for (Int32 i = 0; i < objs.Length; i++)
                objs[i] = a * xcurr[i] + (1 - a) * xbar[i];
            return objs;
        }
        private void LineSearchUpdateVars()
        {
            Double a = _alpha.ToNumber();
            this.SavePrevious();
            this.xbar = a * this.xcurr + (1 - a) * this.xbar;
        }
        private void LineSearch()
        {
            _lsFxn = _currFxn.Subs(this.objective.variables, this.LineSearchSubstitution());
            LineSearchSolver ls = new LineSearchSolver(this, LineSearchMethod.Golden);
            _lstolerance *= _lstoleranceUpdateFactor;
            ls.Tolerance = _lstolerance; // * 1.0 / (1.0 + 0.1 * _iter);
            _prevOptimum2 = _prevOptimum1;
            _prevOptimum1 = _prevOptimum;
            _prevOptimum = _optimum;
            _optimum = ls.Solve(new Object[] { 0.0, 1.0 });
            this.LineSearchUpdateVars();
        }
        private Matrix Gradient()
        {
            return _gradient = Symbol.Eval(Sym.Gradient(_currFxn, this.objective.variables), this.objective.variables, this.xbar);
        }
        private void DifferentiateFunction()
        {
            this.Gradient();
        }
        private bool gradientConverged()
        {
            return _gradient.ElementwiseMult(this.objective.variableMask).Abs() <= _tolerance;
        }
        private bool optimumConverged()
        {
            return !Double.IsNaN(_prevOptimum2)
                && Math.Abs(_optimum - _prevOptimum) <= _tolerance
                && Math.Abs(_prevOptimum - _prevOptimum1) <= _tolerance
                && Math.Abs(_prevOptimum1 - _prevOptimum2) <= _tolerance; 
        }
        private bool xbarConverged()
        {
            return this.xbarprev2 != null
                && (this.xbar - this.xbarprev).Abs() <= _tolerance
                && (this.xbarprev - this.xbarprev1).Abs() <= _tolerance
                && (this.xbarprev1 - this.xbarprev2).Abs() <= _tolerance;
        }
        private void UpdateConvergence()
        {
            Boolean gradConvg = this.gradientConverged();
            Boolean optiConvg = this.optimumConverged(); 
            Boolean xbarConvg = this.xbarConverged();
            Boolean iterMax = _iter >= _maxIters;
            _ConvergedComponents = gradConvg + "," + optiConvg + "," + xbarConvg + "," + iterMax;
            this.objective.IsConverged = gradConvg || optiConvg || xbarConvg || iterMax;
        }
        private void UpdateCosts()
        {
            if (!this.objective.IsConverged)
            {
                Matrix mask = this.objective.variableMask;
                Matrix invMask = 1.0 - mask;
                this.objective.costs = this.objective.costs.ElementwiseMult(invMask) + _gradient.ElementwiseMult(mask);
            }
        }
        public void RunUnconstrainedProblem()
        {
            // get the solver and context
            SolverContext context = new SolverContext();  
            Microsoft.SolverFoundation.Services.Model solver = context.CreateModel(); 

            // goal 
            Symbol goal = _currFxn;

            // add variables 
            List<string> variables = goal.Variables();
            foreach (string variable in variables)
                solver.AddDecision(new Decision(Domain.Real, variable));

            // add a few water constraints
            solver.AddGoal("Unconstrained_Problem", GoalKind.Minimize, goal.ToString());
            Solution solution = context.Solve();

            // write output
            Console.WriteLine("Initial solution:");
            Console.WriteLine("    Quality: " + solution.Quality.ToString());
            Console.WriteLine("    Report:");
            Console.WriteLine(solution.GetReport());

            // assign results to flows
            List<Decision> decs = new List<Decision>(solution.Decisions);
            foreach (Decision dec in decs) 
            {
                int i = Array.FindIndex(this.objective.variables, delegate(String s){ return s == dec.Name; });
                this.xbar[i] = dec.ToDouble();
            }
        }

        private void WriteToLog()
        {
            string line;

            if (_logFile == "") 
                return;

            if (_sw == null)
            {
                try
                {
                    // open the output file if possible
                    _sw = new StreamWriter(_logFile);
                    line = "Tolerance\tIter\tObjective value\talpha\tMax Cost Index\tMax Cost\txbar\tConverged components(gradient,optimum,xbar,itermax)\tObjective String Representation\tGradient Evaluated at xbar";
                    _sw.WriteLine(line);
                }
                catch
                {
                    _logFile = ""; 
                    _sw = null;
                    return;
                }
            }
            Double[] costs = this.objective.costs;
            Double[] nlcosts = this.objective.variableMask.ElementwiseMult(costs).Abs();
            int maxcostindex;
            Double maxcost = Statistics.Max(nlcosts, out maxcostindex);
            line = _tolerance + "\t" + _iter + "\t" + this.Optimum.ToString() + "\t" + _alpha + "\t" 
                + maxcostindex + "\t" + maxcost.ToString() + "\t" + xbar.Transpose().ToString() + "\t" 
                + _ConvergedComponents + "\t" + _currFxn + "\t" + (new Matrix(costs)).Transpose().ToString();  // string.Join("\t", Array.ConvertAll(costs, elem => elem.ToString()));
            _sw.WriteLine(line);
            _sw.Flush();
        }
        private void CloseLog()
        {
            if (_sw != null)
            {
                _sw.Close();
                _sw = null;
            }
        }

        // new threads
        public void SolveOnNewThread()
        {
            this.objective.Initialize();
            _thread = new Thread(() => this.Solve(false));
            _thread.Start();
        }
        public void WaitUntilSolved()
        {
            WaitUntilSolved(0);
        }
        public void WaitUntilSolved(int milliseconds)
        {
            if (_thread != null && _thread.IsAlive)
                if (milliseconds == 0)
                    _thread.Join();
                else
                    _thread.Join(milliseconds); 
        }

        #endregion

        #region Event handlers

        public void StopHandlingModelEvents()
        {
            _handleEvents = false; 
        }
        public void StartHandlingModelEvents()
        {
            _handleEvents = true; 
        }
        void model_Init(object sender, EventArgs e)
        {
            if (_handleEvents)
            {
                this.model.objective.Initialize();
                this.Reset();
            }
            EventHandler eh = this.Init;
            if (eh != null) eh(this, e);
        }
        void model_IterBottom(object sender, EventArgs e)
        {
            if (_handleEvents)
            {
                _iter++;
                this.SolveIter(_iter);
            }
            EventHandler eh = this.IterBottom;
            if (eh != null) eh(this, e);
        }
        void model_Converged(object sender, EventArgs e)
        {
            if (_handleEvents)
            {
                _totalOptimum += _optimum;
                this.Reset();
            }
            EventHandler eh = this.Converged;
            if (eh != null) eh(this, e);
        }

        #endregion

        #region Line search interface components

        /// <summary>Gets and sets the current position in the line search.</summary>
        public Double x { get { return _alpha.Eval(); } set { _alpha = value; } }
        /// <summary>Gets and sets the objective function for </summary>
        public Symbol fxn { get { return _lsFxn; } set { _lsFxn = value; } }

        #endregion
    }
}
